import or1200_definitions
import sys
sys.path.append("..")
import pyrtl
from  pyrtl import GPUSim

width = or1200_definitions.OR1200_OPERAND_WIDTH


class Or1200wbmux(object):
    def __init__(self):
        self.rst = pyrtl.Input(bitwidth=1, name='wbmux_rst')
        self.wb_freeze = pyrtl.Input(bitwidth=1, name='wbmux_wb_freeze')
        self.rfwb_op = pyrtl.Input(bitwidth=width, name='wbmux_rfwb_op')
        self.muxin_a = pyrtl.Input(bitwidth=width, name='wbmux_muxin_a')
        self.muxin_b = pyrtl.Input(bitwidth=width, name='wbmux_muxin_b')
        self.muxin_c = pyrtl.Input(bitwidth=width, name='wbmux_muxin_c')
        self.muxin_d = pyrtl.Input(bitwidth=width, name='wbmux_muxin_d')

        self.muxout_out = pyrtl.Output(bitwidth=width,name='muxout')
        self.muxreg = pyrtl.Output(bitwidth=width,name='muxreg')
        self.muxreg_valid = pyrtl.Output(bitwidth=1,name='muxreg_valid')

        self.muxout = pyrtl.WireVector(bitwidth=width)
        # self.muxreg = pyrtl.Register(bitwidth=width)
        # self.muxreg_valid = pyrtl.Register(bitwidth=1)
        self.muxout_out<<=self.muxout
        # initialize all the parts
        registered_output = Registered_output()
        writeBackMulti = WriteBackMulti()
        
        registered_output.rst <<= self.rst
        registered_output.wb_freeze <<= self.wb_freeze
        registered_output.muxout <<= self.muxout
        registered_output.rfwb_op <<= self.rfwb_op
        self.muxreg <<= registered_output.muxreg
        self.muxreg_valid<<= registered_output.muxreg_valid
        writeBackMulti.muxin_a <<= self.muxin_a
        writeBackMulti.muxin_b <<= self.muxin_b
        writeBackMulti.muxin_c <<= self.muxin_c
        writeBackMulti.muxin_d <<= self.muxin_d
        writeBackMulti.rfwb_op <<= self.rfwb_op
        self.muxout <<= writeBackMulti.muxout


class Or1200wbmux_wire(object):
    def __init__(self):
        self.rst = pyrtl.WireVector(bitwidth=1)
        self.wb_freeze = pyrtl.WireVector(bitwidth=1)
        self.rfwb_op = pyrtl.WireVector(bitwidth=width)
        self.muxin_a = pyrtl.WireVector(bitwidth=width)
        self.muxin_b = pyrtl.WireVector(bitwidth=width)
        self.muxin_c = pyrtl.WireVector(bitwidth=width)
        self.muxin_d = pyrtl.WireVector(bitwidth=width)

        self.muxout_out = pyrtl.WireVector(bitwidth=width)
        self.muxreg = pyrtl.WireVector(bitwidth=width)
        self.muxreg_valid = pyrtl.WireVector(bitwidth=1)

        self.muxout = pyrtl.WireVector(bitwidth=width)
        # self.muxreg = pyrtl.Register(bitwidth=width)
        # self.muxreg_valid = pyrtl.Register(bitwidth=1)
        self.muxout_out <<= self.muxout
        # initialize all the parts
        registered_output = Registered_output()
        writeBackMulti = WriteBackMulti()

        registered_output.rst <<= self.rst
        registered_output.wb_freeze <<= self.wb_freeze
        registered_output.muxout <<= self.muxout
        registered_output.rfwb_op <<= self.rfwb_op
        self.muxreg <<= registered_output.muxreg
        self.muxreg_valid <<= registered_output.muxreg_valid
        writeBackMulti.muxin_a <<= self.muxin_a
        writeBackMulti.muxin_b <<= self.muxin_b
        writeBackMulti.muxin_c <<= self.muxin_c
        writeBackMulti.muxin_d <<= self.muxin_d
        writeBackMulti.rfwb_op <<= self.rfwb_op
        self.muxout <<= writeBackMulti.muxout

# Registered output from the write-back multiplexer
class Registered_output(object):
    def __init__(self):
        self.rst = pyrtl.WireVector(bitwidth=1)
        self.wb_freeze = pyrtl.WireVector(bitwidth=1)
        self.muxout = pyrtl.WireVector(bitwidth=width)
        self.rfwb_op = pyrtl.WireVector(bitwidth=width)

        self.muxreg = pyrtl.WireVector(bitwidth=width)
        self.muxreg_valid = pyrtl.WireVector(bitwidth=1)

        with pyrtl.conditional_assignment:
            with self.rst:
                # muxreg <= #1 32'd0
                self.muxreg |= pyrtl.Const(0, bitwidth=32)
                # muxreg_valid <= #1 1'b0;
                self.muxreg_valid |= pyrtl.Const(0b0,bitwidth=1)
            with ~self.wb_freeze:
                self.muxreg |= self.muxout
                self.muxreg_valid |= self.rfwb_op[0]


#  Write-back multiplexer
class WriteBackMulti(object):
    def __init__(self):
        self.muxin_a = pyrtl.WireVector(bitwidth=width)
        self.muxin_b = pyrtl.WireVector(bitwidth=width)
        self.muxin_c = pyrtl.WireVector(bitwidth=width)
        self.muxin_d = pyrtl.WireVector(bitwidth=width)
        self.rfwb_op = pyrtl.WireVector(bitwidth=width)
        self.muxout = pyrtl.WireVector(bitwidth=width)

        with pyrtl.conditional_assignment:
            with self.rfwb_op[1:or1200_definitions.OR1200_RFWBOP_WIDTH] == pyrtl.Const(0b00, bitwidth=2):
                self.muxout |= self.muxin_a
            with self.rfwb_op[1:or1200_definitions.OR1200_RFWBOP_WIDTH] == pyrtl.Const(0b01, bitwidth=2):
                self.muxout |= self.muxin_b
            with self.rfwb_op[1:or1200_definitions.OR1200_RFWBOP_WIDTH] == pyrtl.Const(0b10, bitwidth=2):
                self.muxout |= self.muxin_c
            with self.rfwb_op[1:or1200_definitions.OR1200_RFWBOP_WIDTH] == pyrtl.Const(0b11, bitwidth=2):
                self.muxout |= self.muxin_d


if __name__ == '__main__':
    or1200wbmux = Or1200wbmux()
    #sim = pyrtl.GPUSim_now.GPUSim('wb')
    sim = pyrtl.GPUSim_now1line.GPUSim('wbb')
    #com = pyrtl.CompiledSimulation()

    input_width = []
    from pyrtl.wire import Input

    IN = pyrtl.working_block().wirevector_subset(Input)
    for inn in IN:
        name = inn.name
        width = inn.bitwidth
        input_width.append([name, width])
    import random

    max_ = 0
    com_time = 0
    ess_time = 0
    num = 1
    for i in range(num):
        print(i)

        inputs_random = {}
        inputs_select = {}
        inputs_old = {}
        inputs_ = {}

        for inp in input_width:
            inputs_random[inp[0]] = []
            inputs_select[inp[0]] = []
            result = random.randint(0, 2 ** inp[1] - 1)
            inputs_old[inp[0]] = result

        for j in range(131072):
            for inp in input_width:
                result = random.randint(0, 2 ** inp[1] - 1)
                inputs_[inp[0]] = result
                inputs_random[inp[0]].append(result)
            if random.random() < 0.2:
                for name in inputs_old:
                    inputs_select[name].append(inputs_old[name])
            else:
                for name in inputs_:
                    inputs_select[name].append(inputs_[name])
                    inputs_old[name] = inputs_[name]
        #time1 = com.step_multiple(inputs_select)

    sim.make_input(inputs_select)